#include "SVMemoryPool.h"

SVMemoryPool::SVMemoryPool(QObject *parent)
    : QObject(parent)
{
}

SVMemoryPool *SVMemoryPool::getInstance()
{
    static SVMemoryPool instance;
    instance.initialize(MEMPOOL_MAX_COUNT, MAT_MAX_SIZE, MAT_MAX_SIZE, CV_8UC3);
    return &instance;
}

void SVMemoryPool::initialize(uint count, uint rows, uint cols, int type)
{
    mMemoryPool.resize(count);
    mMemoryPool.fill(cv::Mat(rows, cols, type));
    for (int i = 0; i < count; ++i) {
        mFreeList.append(i);
    }
}

cv::Mat SVMemoryPool::convertQImageToCVImage(const QImage &image)
{
    const int width = image.width();
    const int height = image.height();
    const size_t step = image.bytesPerLine();
    cv::Mat mat = allocate(cv::Size(width, height));

    switch (image.format()) {
    case QImage::Format_RGB888: {
        cv::Mat temp = cv::Mat(image.height(),
                               image.width(),
                               CV_8UC3,
                               const_cast<uchar *>(image.bits()),
                               image.bytesPerLine());
        cv::cvtColor(temp, mat, cv::COLOR_RGB2BGR);
        break;
    }
    case QImage::Format_Grayscale8: {
        memcpy(mat.data, image.bits(), image.sizeInBytes());
        break;
    }
    default: {
        mat = convertQImageToCVImage(image.convertToFormat(QImage::Format_RGB888));
        break;
    }
    }

    return mat;
}

QImage SVMemoryPool::convertCVImageToQImage(const cv::Mat &mat)
{
    cv::Mat temp = mat.clone();
    QImage image;
    switch (temp.type()) {
    case CV_8UC1: {
        image = QImage(temp.data, temp.cols, temp.rows, temp.step, QImage::Format_Grayscale8);
        break;
    }
    case CV_8UC3: {
        cv::cvtColor(mat, temp, cv::COLOR_BGR2RGB);
        image = QImage(temp.data, temp.cols, temp.rows, temp.step, QImage::Format_RGB888);
        break;
    }
    default:
        break;
    }
    return image;
}

cv::Mat SVMemoryPool::allocate(cv::Size size)
{
    int index = 0;
    if (mFreeList.isEmpty()) {
        mMemoryPool.push_back(cv::Mat(mMemoryPool[0].size(), mMemoryPool[0].type()));
        mFreeList.push_back(mMemoryPool.size() - 1);
    }
    index = mFreeList.takeLast();
    mDataToIndex.insert(mMemoryPool[index].data, index);
    return cv::Mat(mMemoryPool[index], cv::Rect(0, 0, size.width, size.height));
}

void SVMemoryPool::release(cv::Mat &mat)
{
    auto iter = mDataToIndex.find(mat.data);
    if (iter != mDataToIndex.end()) {
        mFreeList.push_back(iter.value());
        mDataToIndex.erase(iter);
    }
}

SVScopedPool::SVScopedPool(SVMemoryPool *pool)
    : mPool(pool)
{
}

SVScopedPool::~SVScopedPool()
{
    for (cv::Mat &mat : mUsedMats) {
        mPool->release(mat);
    }
}

cv::Mat SVScopedPool::convertQImageToCVImage(const QImage &image)
{
    mUsedMats.push_back(mPool->convertQImageToCVImage(image));
    return mUsedMats.back();
}

cv::Mat SVScopedPool::allocate(cv::Size size)
{
    mUsedMats.push_back(mPool->allocate(size));
    return mUsedMats.back();
}
